{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualisation Examples\n",
    "\n",
    "This notebook shows some of the visulisation utility of our toolkit.\n",
    "\n",
    "The core packages for visualisation are:\n",
    "### `rasterization`\n",
    "contains classes for getting visual data as multi-channel tensors and turning them into interpretable RGB images.\n",
    "Every class has at least a `rasterize` method to get the tensor and a `to_rgb` method to convert it into an image.\n",
    "A few examples are:\n",
    "- `BoxRasterizer`: this object renders agents (e.g. vehicles or pedestrians) as oriented 2D boxes\n",
    "- `SatelliteRasterizer`: this object renders an oriented crop from a satellite map\n",
    "\n",
    "### `visualization`\n",
    "contains utilities to draw additional information (e.g. trajectories) onto RGB images. These utilities are commonly used after a `to_rgb` call to add other information to the final visualisation. \n",
    "One example is:\n",
    "- `draw_trajectory`: this function draws 2D trajectories from coordinates and yaws offset on an image\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from l5kit.data import ChunkedStateDataset, LocalDataManager\n",
    "from l5kit.dataset import EgoDataset, AgentDataset\n",
    "\n",
    "from l5kit.rasterization import build_rasterizer\n",
    "from l5kit.configs import load_config_data\n",
    "from l5kit.visualization import draw_trajectory, TARGET_POINTS_COLOR\n",
    "from l5kit.geometry import transform_points\n",
    "from tqdm import tqdm\n",
    "from collections import Counter\n",
    "from l5kit.data import LABELS\n",
    "from prettytable import PrettyTable\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First, let's configure where our data lives!\n",
    "The data is expected to live in a folder that can be configured using the `L5KIT_DATA_FOLDER` env variable. You data folder is expected to contain subfolders for the aerial and semantic maps as well as the scenes (`.zarr` files). \n",
    "In this example, the env variable is set to the local data folder. You should make sure the path points to the correct location for you.\n",
    "\n",
    "We built our code to work with a human-readable `yaml` config. This config file holds much useful information, however, we will only focus on a few functionalities concerning loading and visualization here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set env variable for data\n",
    "os.environ[\"L5KIT_DATA_FOLDER\"] = \"PATH_TO_YOUR_DATA\"\n",
    "# get config\n",
    "cfg = load_config_data(\"./visualisation_config.yaml\")\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### We can look into our current configuration for interesting fields\n",
    "\n",
    "\\- when loaded in python, the `yaml`file is converted into a python `dict`. \n",
    "\n",
    "`raster_params` contains all the information related to the transformation of the 3D world onto an image plane:\n",
    "  - `raster_size`: the image plane size\n",
    "  - `pixel_size`: how many meters correspond to a pixel\n",
    "  - `ego_center`: our raster is centered around an agent, we can move the agent in the image plane with this param\n",
    "  - `map_type`: the rasterizer to be employed. We currently support a satellite-based and a semantic-based one. We will look at the differences further down in this script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'current raster_param: {cfg[\"raster_params\"]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data\n",
    "\n",
    "The same config file is also used to load the data. Every split in the data has its own section, and multiple datasets can be used (as a whole or sliced). In this short example we will only use the first dataset from the `sample` set. You can change this by configuring the 'train_data_loader' variable in the config.\n",
    "\n",
    "You may also have noticed that we're building a `LocalDataManager` object. This will resolve relative paths from the config using the `L5KIT_DATA_FOLDER` env variable we have just set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dm = LocalDataManager()\n",
    "dataset_path = dm.require(cfg[\"train_data_loader\"][\"datasets\"][0][\"key\"])\n",
    "zarr_dataset = ChunkedStateDataset(dataset_path)\n",
    "zarr_dataset.open()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with the raw data\n",
    "\n",
    "`.zarr` files support most of the traditional numpy array operations. In the following cell we iterate over the frames to get a scatter plot of the AV locations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = zarr_dataset.frames\n",
    "coords = np.zeros((len(frames), 2))\n",
    "for idx_coord, idx_data in enumerate(tqdm(range(len(frames)), desc=\"getting centroid to plot trajectory\")):\n",
    "    frame = zarr_dataset.frames[idx_data]\n",
    "    coords[idx_coord] = frame[\"ego_translation\"][:2]\n",
    "\n",
    "\n",
    "plt.scatter(coords[:, 0], coords[:, 1], marker='.')\n",
    "axes = plt.gca()\n",
    "axes.set_xlim([-2500, 1600])\n",
    "axes.set_ylim([-2500, 1600])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another easy thing to try is to get an idea of the agents types distribution. \n",
    "\n",
    "We can get all the agents `label_probabilities` and get the argmax for each raw. because `.zarr` files map to numpy array we can use all the traditional numpy operations and functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agents = zarr_dataset.agents\n",
    "probabilities = agents[\"label_probabilities\"]\n",
    "labels_indexes = np.argmax(probabilities, axis=1)\n",
    "counts = []\n",
    "for idx_label, label in enumerate(LABELS):\n",
    "    counts.append(np.sum(labels_indexes == idx_label))\n",
    "    \n",
    "table = PrettyTable(field_names=[\"label\", \"counts\"])\n",
    "for count, label in zip(counts, LABELS):\n",
    "    table.add_row([label, count])\n",
    "print(table)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with data abstraction\n",
    "\n",
    "Even though it's absolutely fine to work with the raw data, we also provide classes that abstract data access to offer an easier way to generate inputs and targets.\n",
    "\n",
    "### Core Objects\n",
    "Along with the `rasterizer`, our toolkit contains other classes you may want to use while you build your solution. The `dataset` package, for example, already implements `PyTorch` ready datasets, so you can hit the ground running and start coding immediately.\n",
    "\n",
    "### Dataset package\n",
    "We will use two classes from the `dataset` package for this example. Both of them can be iterated and return multi-channel images from the rasterizer along with future trajectories offsets and other information.\n",
    "- `EgoDataset`: this dataset iterates over the AV annotations\n",
    "- `AgentDataset`: this dataset iterates over other agents annotations\n",
    "\n",
    "Both support multi-threading (through PyTorch DataLoader) OOB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rast = build_rasterizer(cfg, dm)\n",
    "dataset = EgoDataset(cfg, zarr_dataset, rast)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What if I want to visualise the Autonomous Vehicle (AV)?\n",
    "\n",
    "Let's get a sample from the dataset and use our `rasterizer` to get an RGB image we can plot. \n",
    "\n",
    "If we want to plot the ground truth trajectory, we can convert the dataset's `target_position` (displacements in meters in world coordinates) into pixel coordinates in the image space, and call our utility function `draw_trajectory` (note that you can use this function for the predicted trajectories, as well)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = dataset[50]\n",
    "\n",
    "im = data[\"image\"].transpose(1, 2, 0)\n",
    "im = dataset.rasterizer.to_rgb(im)\n",
    "target_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\n",
    "draw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n",
    "\n",
    "plt.imshow(im[::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What if I want to change the rasterizer?\n",
    "\n",
    "We can do so easily by building a new rasterizer and new dataset for it. In this example, we change the value to `py_satellite` which renders boxes on an aerial image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg[\"raster_params\"][\"map_type\"] = \"py_satellite\"\n",
    "rast = build_rasterizer(cfg, dm)\n",
    "dataset = EgoDataset(cfg, zarr_dataset, rast)\n",
    "data = dataset[50]\n",
    "\n",
    "im = data[\"image\"].transpose(1, 2, 0)\n",
    "im = dataset.rasterizer.to_rgb(im)\n",
    "target_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\n",
    "draw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n",
    "\n",
    "plt.imshow(im[::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What if I want to visualise an agent?\n",
    "\n",
    "Glad you asked! We can just replace the `EgoDataset` with an `AgentDataset`. Now we're iterating over agents and not the AV anymore, and the first one happens to be the pace car (you will see this one around a lot in the dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = AgentDataset(cfg, zarr_dataset, rast)\n",
    "data = dataset[0]\n",
    "\n",
    "im = data[\"image\"].transpose(1, 2, 0)\n",
    "im = dataset.rasterizer.to_rgb(im)\n",
    "target_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\n",
    "draw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n",
    "\n",
    "plt.imshow(im[::-1])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## System Origin and Orientation\n",
    "\n",
    "At this point you may have noticed that we flip the image on the **Y-axis** before plotting it.\n",
    "\n",
    "When moving from 3D to 2D we stick to a right-hand system, where the origin is in the bottom-left corner with positive x-values going right and positive y-values going up the image plane. The camera is facing down the negative z axis. \n",
    "\n",
    "However, both `opencv` and `pyplot` place the origin in the top-left corner with positive x going right and positive y going down in the image plane. The camera is facing down the positive z-axis.\n",
    "\n",
    "The flip done on the resulting image is for visualisation purposes to accommodate the difference in the two coordinate frames.\n",
    "\n",
    "Further, all our rotations are counter-clockwise for positive value of the angle."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How does an entire scene look like?\n",
    "\n",
    "It's easy to visualise an individual scene using our toolkit. Both `EgoDataset` and `AgentDataset` provide 2 methods for getting interesting indices:\n",
    "- `get_frame_indices` returns the indices for a given frame. For the `EgoDataset` this matches a single observation, while more than one index could be available for the `AgentDataset`, as that given frame may contain more than one valid agent\n",
    "- `get_scene_indices` returns indices for a given scene. For both datasets, these might return more than one index\n",
    "\n",
    "In this example, we visualise the second scene from the ego's point of view:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, clear_output\n",
    "import PIL\n",
    " \n",
    "cfg[\"raster_params\"][\"map_type\"] = \"py_satellite\"\n",
    "rast = build_rasterizer(cfg, dm)\n",
    "dataset = EgoDataset(cfg, zarr_dataset, rast)\n",
    "scene_idx = 1\n",
    "indexes = dataset.get_scene_indices(scene_idx)\n",
    "images = []\n",
    "\n",
    "for idx in indexes:\n",
    "    \n",
    "    data = dataset[idx]\n",
    "    im = data[\"image\"].transpose(1, 2, 0)\n",
    "    im = dataset.rasterizer.to_rgb(im)\n",
    "    target_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\n",
    "    center_in_pixels = np.asarray(cfg[\"raster_params\"][\"ego_center\"]) * cfg[\"raster_params\"][\"raster_size\"]\n",
    "    draw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n",
    "    clear_output(wait=True)\n",
    "    display(PIL.Image.fromarray(im[::-1]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
